package com.baiyang.sharding.page;


import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.*;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.lang.reflect.Method;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Component
@Order(100)
@Intercepts({@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})})
public class PageQueryInterceptor implements Interceptor {
    private static final Map<Class<?>, Method> methodMap = new ConcurrentHashMap<>();
    private static final Set<Class<?>> noOrgClass = new CopyOnWriteArraySet<>();
    private String countSuffix = "_COUNT";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 基础数据准备
        Executor executor = (Executor) invocation.getTarget();
        MappedStatement mappedStatement = (MappedStatement) invocation.getArgs()[0];
        Object parameter = invocation.getArgs()[1];
        RowBounds rowBounds = (RowBounds) invocation.getArgs()[2];
        ResultHandler resultHandler = (ResultHandler) invocation.getArgs()[3];
        ThreadLocal<PageEntity> PAGE_ENTITY = PageHelper.getPageEntity();
        PageEntity pageEntity = PAGE_ENTITY.get();
        // 处理分页请求
        if (Objects.nonNull(pageEntity)) {
            PAGE_ENTITY.remove();
            try {
                // 获取数量
                Long count = count(executor, mappedStatement, parameter, rowBounds);
                if (count > 0) {
                    // 当存在可查找数量时,获取数据
                    PageResult<Object> pageResult = new PageResult<>();
                    BoundSql boundSql = mappedStatement.getBoundSql(parameter);
                    String sql = boundSql.getSql();
                    if (StringUtils.hasLength(pageEntity.getOrderBy())) {
                        sql += " order by " + pageEntity.getOrderBy() + " " + pageEntity.getSort();
                    }
                    sql += " limit " + (pageEntity.getPageOn() - 1) * pageEntity.getPageSize() + ", " + pageEntity.getPageSize();
                    BoundSql newBoundSql = new BoundSql(mappedStatement.getConfiguration(), sql,
                            boundSql.getParameterMappings(), boundSql.getParameterObject());
                    Map<String, Object> additionalParameters = boundSql.getAdditionalParameters();
                    if (additionalParameters != null) {
                        additionalParameters.forEach(newBoundSql::setAdditionalParameter);
                    }
                    MappedStatement newMs = copyFromMappedStatement(mappedStatement, new BoundSqlSqlSource(newBoundSql));
                    invocation.getArgs()[0] = newMs;
                    List<Object> objects = executor.query(newMs, parameter, rowBounds, resultHandler);
                    pageResult.setTotal(count);
                    pageResult.addAll(objects);
                    return pageResult;
                } else {
                    PageResult<Object> pageResult = new PageResult<>();
                    return pageResult;
                }
            } finally {
                PAGE_ENTITY.remove();
            }
        }
        return invocation.proceed();
    }

    private Long count(Executor executor, MappedStatement mappedStatement, Object parameter, RowBounds rowBounds) throws SQLException {
        String mappedId = mappedStatement.getId();
        BoundSql boundSql = mappedStatement.getBoundSql(parameter);
        String sql = boundSql.getSql();
        String[] arraySQL = sql.split("from");
        String countSql = "select count(1) from " + arraySQL[1];
        BoundSql newBoundSql = new BoundSql(mappedStatement.getConfiguration(), countSql,
                boundSql.getParameterMappings(), boundSql.getParameterObject());
        // 适配list
        Map<String, Object> additionalParameters = boundSql.getAdditionalParameters();
        if (additionalParameters != null) {
            additionalParameters.forEach(newBoundSql::setAdditionalParameter);
        }

        List<Object> count = executor.query(newCountMappedStatement(mappedStatement, mappedId + countSuffix, new BoundSqlSqlSource(newBoundSql)), parameter, rowBounds, null);
        return (Long) count.get(0);
    }

    private MappedStatement copyFromMappedStatement(MappedStatement ms, SqlSource newSqlSource) {
        MappedStatement.Builder builder = new MappedStatement.Builder(ms.getConfiguration(), ms.getId(), newSqlSource,
                ms.getSqlCommandType());
        builder.resource(ms.getResource());
        builder.fetchSize(ms.getFetchSize());
        builder.statementType(ms.getStatementType());
        builder.keyGenerator(ms.getKeyGenerator());
        if (ms.getKeyProperties() != null) {
            for (String keyProperty : ms.getKeyProperties()) {
                builder.keyProperty(keyProperty);
            }
        }
        builder.timeout(ms.getTimeout());
        builder.parameterMap(ms.getParameterMap());
        builder.resultMaps(ms.getResultMaps());
        builder.cache(ms.getCache());
        builder.useCache(ms.isUseCache());
        return builder.build();
    }

    private static final List<ResultMapping> EMPTY_RESULTMAPPING = new ArrayList<ResultMapping>(0);

    /**
     * 新建count查询的MappedStatement
     *
     * @param ms
     * @param newMsId
     * @param boundSqlSqlSource
     * @return
     */
    public static MappedStatement newCountMappedStatement(MappedStatement ms, String newMsId, BoundSqlSqlSource boundSqlSqlSource) {
        MappedStatement.Builder builder = new MappedStatement.Builder(ms.getConfiguration(), newMsId, boundSqlSqlSource, ms.getSqlCommandType());
        builder.resource(ms.getResource());
        builder.fetchSize(ms.getFetchSize());
        builder.statementType(ms.getStatementType());
        builder.keyGenerator(ms.getKeyGenerator());
        if (ms.getKeyProperties() != null && ms.getKeyProperties().length != 0) {
            StringBuilder keyProperties = new StringBuilder();
            for (String keyProperty : ms.getKeyProperties()) {
                keyProperties.append(keyProperty).append(",");
            }
            keyProperties.delete(keyProperties.length() - 1, keyProperties.length());
            builder.keyProperty(keyProperties.toString());
        }
        builder.timeout(ms.getTimeout());
        builder.parameterMap(ms.getParameterMap());
        //count查询返回值int
        List<ResultMap> resultMaps = new ArrayList<ResultMap>();
        ResultMap resultMap = new ResultMap.Builder(ms.getConfiguration(), ms.getId(), Long.class, EMPTY_RESULTMAPPING).build();
        resultMaps.add(resultMap);
        builder.resultMaps(resultMaps);
        builder.resultSetType(ms.getResultSetType());
        builder.cache(ms.getCache());
        builder.flushCacheRequired(ms.isFlushCacheRequired());
        builder.useCache(ms.isUseCache());
        return builder.build();
    }


    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {

    }


    public static class BoundSqlSqlSource implements SqlSource {
        BoundSql boundSql;

        public BoundSqlSqlSource(BoundSql boundSql) {
            this.boundSql = boundSql;
        }

        @Override
        public BoundSql getBoundSql(Object parameterObject) {
            return boundSql;
        }
    }



}
